# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: jhkim

import numpy
import pandas


class PnlCompactStat(object):
  def __init__(self, product, filldf, orderdf, symbol_info, squeeze_stat=False):
    self.product = product
    self.filldf = filldf
    # orderdf has very compact columns
    self.orderdf = orderdf
    self.symbol_info = symbol_info
    self.squeeze_stat = squeeze_stat

  @staticmethod
  def from_orderdf(product, orderdf, symbol_info=None, squeeze_stat=False):
    if 'order_id' not in orderdf:
      return None

    orderdf.loc[orderdf['order_id'].dropna().index,
                'order_id'] = orderdf.loc[orderdf['order_id'].dropna().index,
                                          'order_id'].astype(int)
    orderdf = orderdf.rename(columns={'time': 'timestamp'})
    orderdf['timestamp'] = orderdf['timestamp'].astype(int)
    if squeeze_stat:
      orderdf['time_min'] = (orderdf['timestamp'] / 60e9).astype(int)
      orderdf['fill_pq'] = orderdf['fill_price'] * orderdf['fill_qty']
      subdfs = []
      for _, subdf in orderdf.groupby(['type', 'sign', 'post_only', 'fill_type']):
        subdfa = subdf.groupby('time_min').agg({
            'type': 'first',
            'timestamp': 'first',
            'order_price': 'mean',
            'order_qty': 'sum',
            'order_id': 'first',
            'sign': 'first',
            'fill_type': 'first',
            'fill_qty': 'sum',
            'fill_pq': 'sum',
            'post_only': 'first',
        })
        subdfa['fill_price'] = subdfa['fill_pq'] / subdfa['fill_qty']
        subdfs.append(subdfa)
      if len(subdfs) > 0:
        orderdf = pandas.concat(subdfs, axis=0).sort_values('timestamp').reset_index(drop=True)

    if 'fill_qty' not in orderdf:
      orderdf['fill_qty'] = 0
      orderdf['fill_price'] = numpy.nan
      orderdf['fill_type'] = numpy.nan

    optional_columns = ['weight', 'sub_amt_filled']
    for col in optional_columns:
      if not col in orderdf:
        orderdf[col] = 0

    buyps = orderdf['order_price'].copy()
    buyps.loc[orderdf['sign'] != 1] = numpy.nan
    buyps.loc[~orderdf['weight'].isna()] = numpy.nan
    buyps.loc[
        (orderdf['type'] != 'ORDER_ACCEPTED') &
        (orderdf['type'] != 'ORDER_SUBMITTED')] = numpy.nan
    buyps_ts = orderdf['timestamp'].copy()
    buyps = buyps.astype(numpy.float64)
    buyps_ts.loc[numpy.isnan(buyps)] = numpy.nan
    buyps = buyps.fillna(method='ffill')
    sellps = orderdf['order_price'].copy()
    sellps.loc[orderdf['sign'] != -1] = numpy.nan
    sellps.loc[~orderdf['weight'].isna()] = numpy.nan
    sellps.loc[
        (orderdf['type'] != 'ORDER_ACCEPTED') &
        (orderdf['type'] != 'ORDER_SUBMITTED')] = numpy.nan
    sellps_ts = orderdf['timestamp'].copy()
    sellps = sellps.astype(numpy.float64)
    sellps_ts.loc[numpy.isnan(sellps)] = numpy.nan
    sellps = sellps.fillna(method='ffill')
    mtmps = 0.5 * (buyps + sellps)
    # mtmps = mtmps.fillna(method='bfill')

    filldf = orderdf.loc[orderdf['fill_qty'] > 0]

    buysell_ts = pandas.concat(
        [buyps_ts.fillna(method='ffill'), sellps_ts.fillna(method='ffill')],
        axis=1).min(axis=1)

    filldf = filldf.reset_index(drop=True)
    orderdf['mtm_price_from_order'] = mtmps
    orderdf['mtm_price_from_order_ts'] = buysell_ts
    orderdf.rename(columns={'tag': 'order_tag'}, inplace=True)
    orderdf['tag'] = orderdf['type']
    orderdf['tag'] = orderdf['tag'].str.replace("CANCEL_ERROR", "CERROR")
    orderdf['tag'] = orderdf['tag'].str.replace("CANCEL_CONFIRMED", "CAN")

    orderdf = orderdf[[
        'order_id', 'type', 'tag', 'timestamp', 'order_price', 'fill_price', 'order_qty',
        'fill_qty', 'fill_type', 'sign', 'post_only', 'mtm_price_from_order',
        'mtm_price_from_order_ts'
    ] + optional_columns + ["order_tag"] * ("order_tag" in orderdf.columns)
    + ["submit_max_rate_used_ratio"] * ("submit_max_rate_used_ratio" in orderdf.columns)
    + ["cancel_max_rate_used_ratio"] * ("cancel_max_rate_used_ratio" in orderdf.columns)]

    orderdf['order_id'] = orderdf['order_id'].astype(pandas.Int64Dtype())
    if hasattr(symbol_info, "symbol_mtm_time") and symbol_info.symbol_mtm_time > 0:
      orderdf = orderdf.append(dict(
          timestamp=symbol_info.symbol_mtm_time,
          mtm_price_from_order=symbol_info.symbol_mtm_price,
          mtm_price_from_order_ts=symbol_info.symbol_mtm_time), ignore_index=True)
    return PnlCompactStat(
        product=product,
        filldf=filldf,
        orderdf=orderdf,
        symbol_info=symbol_info,
        squeeze_stat=squeeze_stat)
